package org.study.lucene.api.analyzer;

import org.apache.lucene.analysis.*;
import org.apache.lucene.analysis.tokenattributes.*;
import org.wltea.analyzer.core.IKSegmenter;
import org.wltea.analyzer.core.Lexeme;
import org.wltea.analyzer.lucene.IKAnalyzer;

import java.io.IOException;
import java.io.StringReader;
import java.util.ArrayList;
import java.util.List;

/**
 * IKAnalyzer 测试
 *
 * @author Administrator
 * @date 2021-12-29
 */
public class IkAnalyzerTest {
    public static void main(String[] args) throws IOException {
        String english = "Analysis is one of the main causes of slow indexing. Simply put, the more you analyze the slower analyze the indexing (in most cases).";
        String chinese = "厉害了我的哥，中国环保部门发布了治理北京雾霾的的方法！";
        testIKAnalyzer(english);
        testIKAnalyzer(chinese);
        testMyIkAnalyzer1(english);
        testMyIkAnalyzer1(chinese);
        testMyIkAnalyzer2(english);
        testMyIkAnalyzer2(chinese);
    }

    /**
     * 参考 https://blog.csdn.net/napoay/article/details/54864959
     */
    public static void testIKAnalyzer(String text) throws IOException {
        System.out.println("【使用 IKAnalyzer】原语句：" + text);
        try (Analyzer analyzer = new IKAnalyzer(true);
             StringReader reader = new StringReader(text);
             TokenStream tokenStream = analyzer.tokenStream(text, reader)) {
            tokenStream.reset();
            CharTermAttribute teAttribute = tokenStream.getAttribute(CharTermAttribute.class);
            List<String> tokens = new ArrayList<>();
            while (tokenStream.incrementToken()) {
                tokens.add(teAttribute.toString());
            }
            tokenStream.end();
            System.out.println("【使用 IKAnalyzer】分词结果：" + tokens);
        }
    }

    /**
     * IKAnalyzer 细粒度切分。参考 https://blog.csdn.net/u010357298/article/details/80776902
     */
    public static void testMyIkAnalyzer1(String text) throws IOException {
        System.out.println("【使用自定义 IKAnalyzer 实现细粒度切分】原语句：" + text);
        try (Analyzer analyzer = new MyIkAnalyzer();
             TokenStream tokenStream = analyzer.tokenStream("content", text)) {
            tokenStream.reset();
            CharTermAttribute teAttribute = tokenStream.getAttribute(CharTermAttribute.class);
            List<String> tokens = new ArrayList<>();
            while (tokenStream.incrementToken()) {
                tokens.add(teAttribute.toString());
            }
            tokenStream.end();
            System.out.println("【使用自定义 IKAnalyzer 实现细粒度切分】分词结果：" + tokens);
        }
    }

    /**
     * IKAnalyzer 智能切分。参考 https://blog.csdn.net/u010357298/article/details/80776902
     */
    public static void testMyIkAnalyzer2(String text) throws IOException {
        System.out.println("【使用 IKAnalyzer 实现智能切分】原语句：" + text);
        try (Analyzer analyzer = new MyIkAnalyzer(true);
             TokenStream tokenStream = analyzer.tokenStream("content", text)) {
            tokenStream.reset();
            CharTermAttribute teAttribute = tokenStream.getAttribute(CharTermAttribute.class);
            List<String> tokens = new ArrayList<>();
            while (tokenStream.incrementToken()) {
                tokens.add(teAttribute.toString());
            }
            tokenStream.end();
            System.out.println("【使用 IKAnalyzer 实现智能切分】分词结果：" + tokens);
        }
    }
}

/**
 * 重写分析器
 */
class MyIkAnalyzer extends Analyzer {
    private Boolean useSmart;

    public MyIkAnalyzer() {
        this(Boolean.FALSE);
    }

    public MyIkAnalyzer(boolean useSmart) {
        super();
        this.useSmart = useSmart;
    }

    public Boolean isUseSmart() {
        return useSmart;
    }

    public void setUseSmart(Boolean useSmart) {
        this.useSmart = useSmart;
    }

    @Override
    protected TokenStreamComponents createComponents(String fieldName) {
        MyIkTokenizer tk = new MyIkTokenizer(this.useSmart);
        return new TokenStreamComponents(tk);
    }
}

/**
 * 重写分词器
 */
class MyIkTokenizer extends Tokenizer {
    /**
     * IK 分词器实现
     */
    private IKSegmenter ikSegmenter;
    /**
     * 词元文本属性
     */
    private final CharTermAttribute charTermAttribute;
    /**
     * 词元位移属性
     */
    private final OffsetAttribute offsetAttribute;
    /**
     * 词元分类属性(该属性分类参考 org.wltea.analyzer.core.Lexeme 中的分类常量)
     */
    private final TypeAttribute typeAttribute;
    /**
     * 记录最后一个词元的结束位置
     */
    private int endPosition;

    public MyIkTokenizer(boolean useSmart) {
        super();
        offsetAttribute = addAttribute(OffsetAttribute.class);
        charTermAttribute = addAttribute(CharTermAttribute.class);
        typeAttribute = addAttribute(TypeAttribute.class);
        ikSegmenter = new IKSegmenter(input, useSmart);
    }

    @Override
    public boolean incrementToken() throws IOException {
        // 清除所有的词元属性
        clearAttributes();
        Lexeme nextLexeme = ikSegmenter.next();
        if (nextLexeme != null) {
            // 将 Lexeme 转成 Attributes
            // 设置词元文本
            charTermAttribute.append(nextLexeme.getLexemeText());
            // 设置词元长度
            charTermAttribute.setLength(nextLexeme.getLength());
            // 设置词元位移
            offsetAttribute.setOffset(nextLexeme.getBeginPosition(), nextLexeme.getEndPosition());
            // 记录分词的最后位置
            endPosition = nextLexeme.getEndPosition();
            // 记录词元分类
            typeAttribute.setType(nextLexeme.getLexemeTypeString());
            // 返会 true 告知还有下个词元
            return true;
        }
        // 返会 false 告知词元输出完毕
        return false;
    }

    @Override
    public void reset() throws IOException {
        super.reset();
        ikSegmenter.reset(input);
    }

    @Override
    public final void end() {
        int finalOffset = correctOffset(this.endPosition);
        offsetAttribute.setOffset(finalOffset, finalOffset);
    }
}